import arxiv
import asyncio
from typing import List, Dict
from dataclasses import dataclass
from datetime import datetime
from .config import AppConfig

@dataclass
class ArxivPaper:
    id: str
    title: str
    authors: List[str]
    published: datetime
    abstract: str
    categories: List[str]
    pdf_url: str

class ArxivClient:
    def __init__(self, config: AppConfig):
        self.config = config
        self.client = arxiv.Client()

    async def search_papers(self, query: str, max_results: int = 50) -> List[ArxivPaper]:
        """异步搜索arXiv论文"""
        search = arxiv.Search(
            query=query,
            max_results=max_results,
            sort_by=arxiv.SortCriterion.SubmittedDate
        )
        
        papers = []
        async for result in self.client.results(search):
            papers.append(ArxivPaper(
                id=result.get_short_id(),
                title=result.title,
                authors=[a.name for a in result.authors],
                published=result.published,
                abstract=result.summary,
                categories=result.categories,
                pdf_url=result.pdf_url
            ))
            
        return papers

    async def get_trend_data(self, query: str, time_window: str = "month") -> Dict:
        """获取指定时间窗口的趋势数据"""
        papers = await self.search_papers(query)
        
        # 按时间窗口分组统计
        trend_data = {}
        for paper in papers:
            time_key = self._get_time_key(paper.published, time_window)
            trend_data[time_key] = trend_data.get(time_key, 0) + 1
            
        return {
            "query": query,
            "time_window": time_window,
            "trends": [
                {"period": k, "count": v} 
                for k, v in sorted(trend_data.items())
            ],
            "total_papers": len(papers)
        }

    def _get_time_key(self, dt: datetime, window: str) -> str:
        """根据时间窗口生成分组key"""
        if window == "year":
            return str(dt.year)
        elif window == "month":
            return f"{dt.year}-{dt.month:02d}"
        else:  # week
            return f"{dt.year}-W{dt.isocalendar()[1]:02d}"